
import pandas as pd
import numpy as np
import statsmodels.api as sm
from scipy.stats import norm

def rsr_ranking(df: pd.DataFrame, weights: pd.Series, output_path: str, num_levels: int = 4):
    """
    使用秩和秩和相关(RSR)方法对国家进行排名，并通过Probit回归进行分等级。
    
    参数:
        df: 包含国家数据的DataFrame
        weights: 指标权重序列
        output_path: 输出CSV文件路径
        num_levels: 分等级的数量 (3,4,5)
    
    返回:
        None (保存处理后的数据到文件)
    """
    # 确保输入数据类型正确
    if not isinstance(df, pd.DataFrame):
        raise TypeError("df必须是pandas DataFrame")
    if not isinstance(weights, pd.Series):
        raise TypeError("weights必须是pandas Series")
    
    meta_cols = ['序号', '国名En', '国名Ch', '年份']
    data_cols = [col for col in df.columns if col not in meta_cols]

    # 验证数据完整性
    if df[meta_cols].isna().any().any():
        raise ValueError("元数据列包含缺失值")
    if data_cols:
        if df[data_cols].isna().any().any():
            warnings.warn("数据列存在缺失值", RuntimeWarning)
    else:
        raise ValueError("数据列为空")

    # 加权秩次排名
    rank_df = df[data_cols].rank(ascending=True, method='dense')
    total_ranks = rank_df[data_cols].max().max() * len(data_cols)
    df['RSR值'] = (rank_df.sum(axis=1) + 0.5) / total_ranks  # 加1/2避免极端值
    df['RSR排名'] = df['RSR值'].rank(method='min', ascending=False)

    # 年度内排名
    df['年度排名'] = df.groupby('年份')['RSR排名'].rank(method='dense')

    # Probit回归处理
    df_clean = df.copy()
    min_val, max_val = df_clean['RSR值'].min(), df_clean['RSR值'].max()
    
    # 转换RSR值到正态分布
    df_clean['Probit值'] = norm.ppf(df_clean['RSR值'], epsilon, 1-epsilon)
    df_clean['Probit值'] = pd.Series(min(df_clean['Probit值'].values, key=lambda x: abs(x)))
    df_clean['Probit值'].replace([np.inf, -np.inf], np.nan, inplace=True)
    df_clean['Probit值'].fillna(df_clean['Probit值'].mean(), inplace=True)

    model = sm.OLS(df_clean['RSR值'], sm.add_constant(df_clean['Probit值'])).fit()
    df_clean['RSR拟合值'] = model.predict(sm.add_constant(df_clean['Probit值']))

    # 确定分界点
    valid_levels = {3, 4, 5}
    if num_levels not in valid_levels:
        raise ValueError(f"档次数必须是3,4或5，当前为: {num_levels}")
    
    probabilities = np.linspace(0, 1, num_levels+2)[1:-1]
    cut_points = model.params['const'] + model.params['Probit值'] * norm.ppf(probabilities)
    
    # 创建分档
    bins = [-np.inf] + list(cut_points) + [np.inf]
    df_clean['分档等级'] = pd.cut(
        df_clean['RSR拟合值'], 
        bins=bins, 
        labels=range(1, num_levels+1),
        ordered=True
    )

    # 输出结果
    df_clean[['序号', '国名En', '国名Ch', '年份'] + data_cols + ['RSR值', 'RSR排名', '年度排名', '分档等级']].to_csv(
        output_path, 
        index=False, 
        encoding='utf-8-sigma'
    )
    print(f"📄 RSR 分档排序结果已保存到：{output_path}")
